import sys
import os
res=os.path.abspath(__file__) #获取当前文件的绝对路径
base_path=os.path.dirname(os.path.dirname(res)) #获取当前文件的上两级目录
sys.path.insert(1, base_path)

base_path_pre=os.path.dirname(base_path)
sys.path.insert(1, base_path_pre)

import argparse
import copy
import os
import os.path as osp
import time
import warnings
from timm.utils import AverageMeter
import datetime

os.environ['CUDA_VISIBLE_DEVICES']="3"

import mmcv
import torch
import torchvision
import torch.distributed as dist
from mmcv import Config, DictAction
from mmcv.runner import get_dist_info, init_dist
from mmcv.utils import get_git_hash

from mmdet import __version__
from mmdet.apis import init_random_seed, set_random_seed, train_detector
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.utils import (collect_env, get_device, get_root_logger,
                         replace_cfg_vals, setup_multi_processes,
                         update_data_root)
import mmcv_custom  # noqa: F401,F403
import mmdet_custom  # noqa: F401,F403

from mmdet_custom.models import DINOHead
from tools.utils import *
import numpy as np
from mmdet.apis import multi_gpu_test, single_gpu_test
from allinone_new.allinone import ALLINONE
from mmdet.datasets.coco import CocoDataset
from mmdet.datasets.pipelines import LoadImageFromFile
from mmdet.datasets.pipelines import RandomFlip, LoadImageFromFile, LoadAnnotations, AutoAugment, Normalize, Pad, DefaultFormatBundle, Collect
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

from mmcv.runner import BaseRunner
from mmcv.runner import EpochBasedRunner
from mmcv.runner.hooks import StepLrUpdaterHook
from torch.optim.lr_scheduler import StepLR
import torch.distributed as dist
import numpy as np
from tools.utils import SubsetRandomSampler
    
from torch.utils.data import DataLoader
from functools import partial
from mmcv.parallel import collate
from mmcv.runner import (get_dist_info, init_dist, wrap_fp16_model)
    
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
# from mmcv.runner import load_checkpoint
from mmdet.core import DistEvalHook, EvalHook, build_optimizer
from tools.utils import auto_scale_lr
import mmcv
# OPTIMIZERS就是定义在mmcv/runner/optimier/builder.py中，注册了pytorch中的优化器
from torch.optim.adamw import AdamW
import  mmcv.runner.optimizer.builder
from tools.utils import build_scheduler
from mmdet.datasets.pipelines import Collect
from mmcv.image import tensor2imgs
from mmdet.core import encode_mask_results

def parse_args():
    parser = argparse.ArgumentParser(description='Train a all task in on model')
    parser.add_argument('--config', help='train config file path', default="/mnt/pde/algorithm/user/qxu/project/InternImage/allinone_new/configs/coco/all_in_one.py")
    group_gpus = parser.add_mutually_exclusive_group()
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='(Deprecated, please use --gpu-id) number of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='(Deprecated, please use --gpu-id) ids of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument('--gpu-id',
                            type=int,
                            default=0,
                            help='id of gpu to use '
                            '(only applicable to non-distributed training)')
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    parser.add_argument(
        '--diff-seed',
        action='store_true',
        help='Whether or not set different seeds for different ranks')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument(
        '--options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file (deprecate), '
        'change to --cfg-options instead.')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument('--launcher',
                        choices=['none', 'pytorch', 'slurm', 'mpi'],
                        default='none',
                        help='job launcher')
    parser.add_argument('--local_rank', type=int)
    parser.add_argument('--auto-scale-lr',
                        action='store_true',
                        help='enable automatically scaling LR.')
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    if args.options and args.cfg_options:
        raise ValueError(
            '--options and --cfg-options cannot be both '
            'specified, --options is deprecated in favor of --cfg-options')
    if args.options:
        warnings.warn('--options is deprecated in favor of --cfg-options')
        args.cfg_options = args.options

    return args


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)

    # replace the ${key} with the value of cfg.key
    cfg = replace_cfg_vals(cfg)
    # update data root according to MMDET_DATASETS
    update_data_root(cfg)

    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    # # set multi-process settings
    # setup_multi_processes(cfg)
    

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # work_dir is determined in this priority: CLI > segment in file > filename
    if cfg.work_dir is not None:
        # update configs according to CLI cfg if cfg.work_dir is not None
        cfg.work_dir = cfg.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(cfg.config))[0])

    if cfg.resume_from is not None:
        cfg.resume_from = cfg.resume_from
    if args.gpus is not None:
        cfg.gpu_ids = range(1)
        warnings.warn('`--gpus` is deprecated because we only support '
                      'single GPU mode in non-distributed training. '
                      'Use `gpus=1` now.')
    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids[0:1]
        warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
                      'Because we only support single GPU mode in '
                      'non-distributed training. Use the first GPU '
                      'in `gpu_ids` now.')
    if args.gpus is None and args.gpu_ids is None:
        cfg.gpu_ids = [args.gpu_id]

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
        # re-set gpu_ids with distributed training mode
        _, world_size = get_dist_info()
        cfg.gpu_ids = range(world_size)

    # create work_dir
    os.makedirs(cfg.work_dir, exist_ok=True)
    # dump configced before
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init the logger before other stepsment '
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    meta = dict()
    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    meta['env_info'] = env_info
    meta['config'] = cfg.pretty_text
    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config:\n{cfg.pretty_text}')

    cfg.device = get_device()
    # set random seeds
    seed = init_random_seed(args.seed, device=cfg.device)
    seed = seed + dist.get_rank() if args.diff_seed else seed
    logger.info(f'Set random seed to {seed}, '
                f'deterministic: {args.deterministic}')
    set_random_seed(seed, deterministic=args.deterministic)
    cfg.seed = seed
    meta['seed'] = seed
    meta['exp_name'] = osp.basename(args.config)
    

    ## model define
    model = ALLINONE(cfg) # model need init TODO

    # task = ['cls', 'detection', 'segmention']
    for task in cfg.task:
        if task == 'cls':
            #prepare data loaders
            data_transform = transforms.Compose([
                transforms.Resize(size=(224, 224)),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225])
                ])
            train_dataset = torchvision.datasets.ImageFolder(root=cfg.cls_config.data.data_path+'/train/', transform=data_transform)
            train_dataset_loader = DataLoader(train_dataset, batch_size=cfg.cls_config.data.batch_size, shuffle=True, num_workers=cfg.cls_config.data.num_workers)

            val_dataset = torchvision.datasets.ImageFolder(root=cfg.cls_config.data.data_path+'/val/', transform=data_transform)
            val_dataset_loader = DataLoader(val_dataset, batch_size=cfg.cls_config.data.batch_size, shuffle=False, num_workers=cfg.cls_config.data.num_workers)
            
            # linear scale the learning rate according to total batch size, may not be optimal
            rank, world_size = get_dist_info()
            linear_scaled_lr = cfg.cls_config.train.base_lr *  cfg.cls_config.data.batch_size * world_size / 512.0
            linear_scaled_warmup_lr = cfg.cls_config.train.warmup_lr * cfg.cls_config.data.batch_size * world_size / 512.0
            linear_scaled_min_lr = cfg.cls_config.train.min_lr * cfg.cls_config.data.batch_size * world_size / 512.0
            
            cfg.cls_config.train.base_lr = linear_scaled_lr
            cfg.cls_config.train.warmup_lr = linear_scaled_warmup_lr
            cfg.cls_config.train.min_lr = linear_scaled_min_lr
            
            os.makedirs(cfg.work_dir, exist_ok=True)
            
            if rank == 0:
                path = os.path.join(cfg.work_dir, "cls_config.json")
                with open(path, "w") as f:
                    f.write(cfg.dump())
                logger.info(f"Full config saved to {path}")
            
            logger.info(cfg.dump())
            
            model.cuda()
            logger.info(str(model))
            
            # build optimizer
            optimizer = build_cls_optimizer(cfg, model)
            
            # put model on gpu
            if not distributed:
                model = model.cuda()
                model_without_ddp = model
            else:
                model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[torch.cuda.current_device()], broadcast_buffers=False)
                model_without_ddp = model.module
            n_parameters = sum(p.numel()for p in model.parameters() if p.requires_grad)
            logger.info(f"number of params: {n_parameters}")
            if hasattr(model_without_ddp, 'flops'):
                flops = model_without_ddp.flops()
                logger.info(f"number of GFLOPs: {flops / 1e9}")
            
            # build learning rate scheduler
            lr_scheduler = build_scheduler(config=cfg.cls_config, optimizer=optimizer, n_iter_per_epoch=len(train_dataset_loader)) if not cfg.cls_config.eval_model else None
            
            # set auto resume
            resume_file = None
            if cfg.cls_config.train.auto_resume:
                resume_file = auto_resume_helper(cfg.work_dir)
                if resume_file:
                    if cfg.cls_config.train.model_resume:
                        logger.warning(f"auto-resume changing resume file from {cfg.cls_config.train.model_resume} to {resume_file}")
                    cfg.cls_config.train.model_resume = resume_file
                    logger.info(f'auto resuming from {resume_file}')
                else:
                    logger.info(f'no checkpoint found in {cfg.work_dir}, ignoring auto resume')

            if cfg.cls_config.train.model_resume:
                loss_record = load_checkpoint(
                    cfg.cls_config, model_without_ddp, optimizer, lr_scheduler, logger, task="cls")
                # val_cls(cfg.cls_config, val_dataset_loader, model_without_ddp, logger,  epoch=None)
            
            # train
            logger.info("Start training cls")
            start_time = time.time()
            cfg = cfg.cls_config
            for epoch in range(cfg.train.start_epoch, cfg.train.epochs):
                # train_dataset_loader.sampler.set_epoch(epoch)
                train_loss = train_one_epoch_cls(cfg, model, train_dataset_loader, optimizer, epoch, lr_scheduler, logger)

                if distributed:
                    if dist.get_rank() == 0 and (epoch % cfg.SAVE_FREQ == 0 or epoch == (cfg.TRAIN.EPOCHS - 1)):
                        save_checkpoint(cfg, epoch, model_without_ddp, 
                                    train_loss, optimizer, lr_scheduler, logger, task_name="cls")
                else:
                    save_checkpoint(cfg, epoch, model_without_ddp,
                                    train_loss, optimizer, lr_scheduler, logger, task_name="cls")
                
                val_cls(cfg, val_dataset_loader, model, logger,  epoch=epoch)
                
        if task == 'detection':
            loss_record = 1000000000000
            # detection dateset build
            # use CocoDataset
            datasets = [build_dataset(cfg.detection_config.data.train)]
            
            # use pipeline process data
            if len(cfg.detection_config.data.workflow) == 2:
                val_dataset = copy.deepcopy(cfg.detection_config.data.val)
                datasets.append(build_dataset(val_dataset))
            
            model.CLASSES = datasets[0].CLASSES

            ## train process
            # step 1 build dataloader
            # prepare data loaders
            dataset = datasets if isinstance(datasets, (list, tuple)) else [datasets]
            dataset_train = dataset[0]
            dataset_val = dataset[1]

            global_rank, num_tasks = get_dist_info()
            
            sampler_train = torch.utils.data.DistributedSampler(
                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True)
            
            sampler_val = torch.utils.data.DistributedSampler(
                dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
            
            data_loader_train = DataLoader(
                dataset_train, sampler=sampler_train,
                batch_size=cfg.detection_config.data.samples_per_gpu,
                num_workers=cfg.detection_config.data.workers_per_gpu,
                collate_fn=partial(collate, samples_per_gpu=cfg.detection_config.data.samples_per_gpu),
                pin_memory=True,
                drop_last=True,
            )
            
            data_loader_val = DataLoader(
                dataset_val, sampler=sampler_val,
                batch_size=1,
                num_workers=1,
                shuffle=False,
                collate_fn=partial(collate, samples_per_gpu=1),
                pin_memory=False,
                drop_last=False
            )
            
            if not distributed:
                model = MMDataParallel(model, device_ids=cfg.gpu_ids)
            else:
                model = MMDistributedDataParallel(
                    model.cuda(),
                    device_ids=[torch.cuda.current_device()],
                    broadcast_buffers=False)
                
            model_without_ddp = model.module
            n_parameters = sum(p.numel()for p in model.parameters() if p.requires_grad)
            logger.info(f"number of params: {n_parameters}")
            
            
            # linear scale the learning rate according to total batch size, may not be optimal
            rank, world_size = get_dist_info()
            
            linear_scaled_lr = cfg.detection_config.optimizer.lr *  cfg.detection_config.data.samples_per_gpu * world_size / cfg.detection_config.data.base_batch_size
            linear_scaled_warmup_lr = cfg.detection_config.train.warmup_lr * cfg.detection_config.data.samples_per_gpu * world_size / cfg.detection_config.data.base_batch_size
            linear_scaled_min_lr = cfg.detection_config.train.min_lr * cfg.detection_config.data.samples_per_gpu * world_size / cfg.detection_config.data.base_batch_size
            
            cfg.detection_config.train.base_lr = linear_scaled_lr
            cfg.detection_config.train.warmup_lr = linear_scaled_warmup_lr
            cfg.detection_config.train.min_lr = linear_scaled_min_lr
            
            ## step 3 build optimizer
            optimizer = build_optimizer(model_without_ddp, cfg.detection_config.optimizer)
            # print(optimizer)
            
            
            ## step 4 later lr
            lr_scheduler = build_scheduler(config=cfg.detection_config, optimizer=optimizer, n_iter_per_epoch=len(data_loader_train)) # need check again
            
            ## load ckpt for result_path
            # resume for ckpt
            # cfg.defrost()
            resume_file = None
            if cfg.detection_config.train.auto_resume:
                resume_file = auto_resume_helper(cfg.work_dir)
                if resume_file:
                    if cfg.detection_config.train.model_resume:
                        logger.warning(f"auto-resume changing resume file from {cfg.detection_config.train.model_resume} to {resume_file}")
                    cfg.detection_config.train.model_resume = resume_file
                    logger.info(f'auto resuming from {resume_file}')
                else:
                    logger.info(f'no checkpoint found in {cfg.work_dir}, ignoring auto resume')

            if cfg.detection_config.train.model_resume:
                loss_record = load_checkpoint(
                    cfg.detection_config, model_without_ddp, optimizer, lr_scheduler, logger, task="detection")
                
                # val_detection(model, data_loader_val, cfg.detection_config, logger)
            
            
            ## step 5 train process
            cfg = cfg.detection_config
            for epoch in range(cfg.train.start_epoch, cfg.train.epochs):
                data_loader_train.sampler.set_epoch(epoch)
                train_loss = train_one_epoch_detection(cfg, model, data_loader_train, optimizer, epoch, lr_scheduler, logger)

                if distributed:
                    if dist.get_rank() == 0 and (epoch % cfg.SAVE_FREQ == 0 or epoch == (cfg.TRAIN.EPOCHS - 1)):
                        save_checkpoint(cfg, epoch, model_without_ddp, 
                                    train_loss, optimizer, lr_scheduler, logger, task_name="detection")
                else:
                    save_checkpoint(cfg, epoch, model_without_ddp,
                                    train_loss, optimizer, lr_scheduler, logger, task_name="detection")
                
                val_detection(model, data_loader_val, cfg, logger, epoch)
                   

def val_detection(model, data_loader_val, cfg, logger, epoch=None):
        ## do validation
        
        outputs = detection_single_gpu_test(model, data_loader_val, cfg.test.show, cfg.test.show_dir)
        rank, _ = get_dist_info()
        if cfg.output is not None and rank == 0:
            os.makedirs(cfg.output, exist_ok=True)
            json_file = osp.join(cfg.output, f'eval_detection_epoch_{epoch}.json')
        if rank == 0:
            if cfg.test.out:
                print(f'\nwriting results to {cfg.test.out}')
                mmcv.dump(outputs, cfg.test.out)
            kwargs = {}
            if cfg.test.eval:
                eval_kwargs = cfg.test.get('evaluation', {}).copy()
                for key in [
                        'interval', 'tmpdir', 'start', 'gpu_collect', 'save_best',
                        'rule', 'dynamic_intervals'
                ]:
                    eval_kwargs.pop(key, None)
                eval_kwargs.update(dict(metric=cfg.test.eval, **kwargs))
                metric = data_loader_val.dataset.evaluate(outputs, **eval_kwargs)
                print(metric)
                metric_dict = dict(config=cfg, metric=metric)
                if cfg.output is not None and rank == 0:
                    mmcv.dump(metric_dict, json_file)
            
        logger.info(f'++++++++++++++++++++++eval epoch {epoch} as fellow+++++++++++++++++++++++')
        logger.info(f'{metric}')
        logger.info(f'++++++++++++++++++++++++++++++++++++++++++++')

def detection_single_gpu_test(model,
                    data_loader,
                    show=False,
                    out_dir=None,
                    show_score_thr=0.3):
    model.eval()
    results = []
    dataset = data_loader.dataset
    PALETTE = getattr(dataset, 'PALETTE', None)
    prog_bar = mmcv.ProgressBar(len(dataset))
    # print("self.img_ids: ", dataset.img_ids[:5])
    for i, data in enumerate(data_loader):
        # print(data['img_metas'][0].data[0][0]['ori_filename'])
        with torch.no_grad():
            result = model(task="detection", return_loss=False, rescale=True, **data)

        batch_size = len(result)
        if show or out_dir:
            if batch_size == 1 and isinstance(data['img'][0], torch.Tensor):
                img_tensor = data['img'][0]
            else:
                img_tensor = data['img'][0].data[0]
            img_metas = data['img_metas'][0].data[0]
            imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg'])
            assert len(imgs) == len(img_metas)

            for i, (img, img_meta) in enumerate(zip(imgs, img_metas)):
                h, w, _ = img_meta['img_shape']
                img_show = img[:h, :w, :]

                ori_h, ori_w = img_meta['ori_shape'][:-1]
                img_show = mmcv.imresize(img_show, (ori_w, ori_h))

                if out_dir:
                    out_file = osp.join(out_dir, img_meta['ori_filename'])
                else:
                    out_file = None

                show_result(
                    img_show,
                    result[i],
                    bbox_color=PALETTE,
                    text_color=PALETTE,
                    mask_color=PALETTE,
                    show=show,
                    out_file=out_file,
                    score_thr=show_score_thr,
                    classes=data_loader.dataset.CLASSES)

        # encode mask results
        if isinstance(result[0], tuple):
            result = [(bbox_results, encode_mask_results(mask_results))
                      for bbox_results, mask_results in result]
        # This logic is only used in panoptic segmentation test.
        elif isinstance(result[0], dict) and 'ins_results' in result[0]:
            for j in range(len(result)):
                bbox_results, mask_results = result[j]['ins_results']
                result[j]['ins_results'] = (bbox_results,
                                            encode_mask_results(mask_results))

        results.extend(result)

        for _ in range(batch_size):
            prog_bar.update()
    return results   
            
def train_one_epoch_detection(config, model, data_loader, optimizer, epoch, lr_scheduler, logger):
    model.train()
    optimizer.zero_grad()
    # record
    
    num_steps = len(data_loader)
    batch_time = AverageMeter()
    loss_meter = AverageMeter()
    
    import datetime 
    start = time.time()
    end = time.time()

    for idx, batch in enumerate(data_loader):  # TODO
        # if idx == 1:
        #     break
        # print(idx, batch)

        loss, out, log_vars = model(task='detection', **batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        lr_scheduler.step_update(epoch * num_steps + idx)

        torch.cuda.synchronize()
        
        loss_meter.update(loss.item(), batch['img']._data[0].size(0))
        batch_time.update(time.time() - end)
        end = time.time()

        if idx % config.print_freq == 0: # use config
            lr = optimizer.param_groups[0]['lr']
            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
            etas = batch_time.avg * (num_steps - idx)
            logger.info(
                f'Train: [{epoch}/{config.train.epochs}][{idx}/{num_steps}]\t'
                f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.9f}\t'
                f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t'
                f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
                f'loss all {log_vars}\t'
                f'mem {memory_used:.0f}MB')
    epoch_time = time.time() - start
    logger.info(f"Detection training epoch {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}")
    return np.round(loss_meter.avg, 2)


# fellow used for cls
def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    maxk = min(max(topk), output.size()[1])
    batch_size = target.size(0)
    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.reshape(1, -1).expand_as(pred))
    return [correct[:min(k, maxk)].reshape(-1).float().sum(0) * 100. / batch_size for k in topk]

@torch.no_grad()
def val_cls(config, data_loader, model, logger,  epoch=None):
    model.eval()

    batch_time = AverageMeter()
    loss_meter = AverageMeter()
    acc1_meter = AverageMeter()
    acc5_meter = AverageMeter()

    end = time.time()
    for idx, (images, target) in enumerate(data_loader):
        images = images.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)
        loss, output = model(img=images, gt_labels=target, task="classification")

        # convert 22k to 1k to evaluate
        if output.size(-1) == 21841:
            convert_file = './meta_data/map22kto1k.txt'
            with open(convert_file, 'r') as f:
                convert_list = [int(line) for line in f.readlines()]
            output = output[:, convert_list]

        # measure accuracy and record loss
        # loss = criterion(output, target)
        acc1, acc5 = accuracy(output, target, topk=(1, 5))


        loss_meter.update(loss, target.size(0))
        acc1_meter.update(acc1.item(), target.size(0))
        acc5_meter.update(acc5.item(), target.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if idx % config.print_freq == 0:
            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
            logger.info(f'Test: [{idx}/{len(data_loader)}]\t'
                        f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                        f'Loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
                        f'Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t'
                        f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t'
                        f'Mem {memory_used:.0f}MB')
    if epoch is not None:
        logger.info(
            f'[Epoch:{epoch}] * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}'
        )
    else:
        logger.info(
            f' * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}')

    return acc1_meter.avg, acc5_meter.avg, loss_meter.avg

def train_one_epoch_cls(config, model, data_loader, optimizer, epoch, lr_scheduler, logger):
    model.train()
    optimizer.zero_grad()

    num_steps = len(data_loader)
    batch_time = AverageMeter()
    model_time = AverageMeter()
    loss_meter = AverageMeter()

    start = time.time()
    end = time.time()
    for idx, (samples, targets) in enumerate(data_loader):
        iter_begin_time = time.time()
        samples = samples.cuda(non_blocking=True)
        targets = targets.cuda(non_blocking=True)
        
        loss, output = model(img=samples, gt_labels=targets, task="classification")
        loss.backward()
        optimizer.step()
        lr_scheduler.step_update(epoch * num_steps + idx)
        
        torch.cuda.synchronize()

        loss_meter.update(loss.item(), targets.size(0))
        batch_time.update(time.time() - end)
        model_time.update(time.time() - iter_begin_time)
        end = time.time()

        if idx % config.print_freq == 0:
            lr = optimizer.param_groups[0]['lr']
            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
            etas = batch_time.avg * (num_steps - idx)
            logger.info(
                f'Train: [{epoch}/{config.train.epochs}][{idx}/{num_steps}]\t'
                f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.9f}\t'
                f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t'
                f'model_time {model_time.val:.4f} ({model_time.avg:.4f})\t'
                f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
                f'mem {memory_used:.0f}MB')
    epoch_time = time.time() - start
    logger.info(
        f"Classification training epoch {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}"
    )
    return np.round(loss_meter.avg, 2)

if __name__ == '__main__':
    main()